//
// Created by DELL on 2023/3/27.
//

#include "MemoryPool.h"

void* MemoryPool::allocate(std::size_t size)
{
    Block* block{};
    size = ((size + (MinBlockSize - 1) ) / MinBlockSize) * MinBlockSize;

    for (int i = (int)freeBlocks.size() - 1; i >= 0; --i)
    {
        const auto& b = freeBlocks[i];
        if (b->size >= size)
        {
            block = b;
            freeBlocks.erase(freeBlocks.begin() + i);
            break ;
        }
    }

    if (block == nullptr)
    {
        auto chunk = this->createChunk(size);
        block = &chunk->blocks.front();
    }
    void* pAllocated = block->start;

    if (block->size == size)
    {
        block->isAllocated = true;
    }
    else
    {
        Block remainBlock;
        remainBlock.chunk = block->chunk;
        remainBlock.start = reinterpret_cast<uint8_t*>(block->start) + size;
        remainBlock.size = block->size - size;
        remainBlock.isAllocated = false;
        
        block->size = size;
        block->isAllocated = true;

        auto& blocks = block->chunk->blocks;
        auto it = blocks.begin();
        for (;it != blocks.end(); ++it)
        {
            if (block == &(*it))
            {
                ++it;
                it = blocks.insert(it, remainBlock);
                break ;
            }
        }

        this->addFreeBlock(&(*it));
    }
    return pAllocated;
}

void MemoryPool::free(void* p)
{
    Chunk* chunk{};
    for (auto& c : this->chunks)
    {
        if (p >= c.p && p < reinterpret_cast<uint8_t*>(c.p) + c.size)
        {
            chunk = &c;
            break ;
        }
    }
    if (chunk == nullptr)
    {
        return ;
    }

    auto it = chunk->blocks.begin();
    for (; it != chunk->blocks.end(); ++it)
    {
        if (p >= it->start && p < reinterpret_cast<uint8_t*>(it->start) + it->size)
        {
            break ;
        }
    }

    if (it == chunk->blocks.end())
    {
        return ;
    }

    it->isAllocated = false;
    auto next = it;
    next++;
    if (next != chunk->blocks.end() && !next->isAllocated)
    {
        it->size += next->size;
        this->removeFreeBlock(&(*next));
        chunk->blocks.erase(next);
    }
    if (it != chunk->blocks.begin())
    {
        auto preview = it;
        --preview;
        if (!preview->isAllocated)
        {
            it->size += preview->size;
            this->removeFreeBlock(&(*preview));
            chunk->blocks.erase(preview);
        }
    }
    this->addFreeBlock(&(*it));
}

MemoryPool::Chunk* MemoryPool::createChunk(uint32_t size)
{
    if (size <= MinChunkSize)
    {
        size = MinChunkSize;
    }
    else if (size <= MaxChunkSize)
    {
        for (int i = 21; i <= 25; ++i)
        {
            if (size <= (1 << i))
            {
                size = (1 << i);
                break ;
            }
        }
    }
    else
    {
        size = ((size / MaxChunkSize) + 1) * MaxChunkSize;
    }

    auto& chunk = this->chunks.emplace_back(operator new(size, std::align_val_t{MinBlockSize}), size);
    chunk.blocks.emplace_back(&chunk, chunk.p, size, false);
    return &chunk;
}

void MemoryPool::addFreeBlock(Block* block)
{
    this->freeBlocks.push_back(block);
    int i = (int)freeBlocks.size() - 2;
    for ( ;i >= 0; --i)
    {
        if (freeBlocks[i]->size < block->size)
        {
            freeBlocks[i + 1] = freeBlocks[i];
        }
        else
        {
            freeBlocks[i + 1] = block;
            break ;
        }
    }
    if (i < 0)
    {
        freeBlocks[0] = block;
    }
}

void MemoryPool::removeFreeBlock(Block* block)
{
    auto it = this->freeBlocks.begin();
    while (it != this->freeBlocks.end())
    {
        if ((*it) == block)
        {
            this->freeBlocks.erase(it);
            break ;
        }
        ++it;
    }
}

